#!/usr/bin/env bash
set -euo pipefail

# One-click setup and run script for Linux
# - Creates a Python venv
# - Installs compatible CPU (default) or CUDA PyTorch + PyG
# - Installs FastAPI & Uvicorn
# - Checks required assets
# - Starts the FastAPI server

PROJECT_ROOT="$(cd "$(dirname "$0")"/.. && pwd)"
cd "$PROJECT_ROOT"

PYTHON_BIN="python3"
VENV_DIR=".venv"
PORT="8000"
HOST="0.0.0.0"
TORCH_VER="2.3.1"
TV_VER="0.18.1"
TA_VER="2.3.1"
CUDA_VARIANT="cpu"  # set to cu121 for CUDA 12.1, cu118 for CUDA 11.8

usage() {
  echo "Usage: $0 [--cuda cu121|cu118|cpu] [--port PORT] [--host HOST]"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --cuda)
      CUDA_VARIANT="$2"; shift 2;;
    --port)
      PORT="$2"; shift 2;;
    --host)
      HOST="$2"; shift 2;;
    -h|--help)
      usage; exit 0;;
    *)
      echo "Unknown arg: $1"; usage; exit 1;;
  esac
done

echo "[1/6] Creating venv at $VENV_DIR"
if [[ ! -d "$VENV_DIR" ]]; then
  "$PYTHON_BIN" -m venv "$VENV_DIR"
fi
source "$VENV_DIR/bin/activate"

echo "[2/6] Upgrading pip"
python -m pip install --upgrade pip wheel setuptools

echo "[3/6] Installing PyTorch ($CUDA_VARIANT) $TORCH_VER and dependencies"
if [[ "$CUDA_VARIANT" == "cpu" ]]; then
  TORCH_INDEX_URL="https://download.pytorch.org/whl/cpu"
  PYG_WHL_URL="https://data.pyg.org/whl/torch-${TORCH_VER}+cpu.html"
elif [[ "$CUDA_VARIANT" == "cu121" ]]; then
  TORCH_INDEX_URL="https://download.pytorch.org/whl/cu121"
  PYG_WHL_URL="https://data.pyg.org/whl/torch-${TORCH_VER}+cu121.html"
elif [[ "$CUDA_VARIANT" == "cu118" ]]; then
  TORCH_INDEX_URL="https://download.pytorch.org/whl/cu118"
  PYG_WHL_URL="https://data.pyg.org/whl/torch-${TORCH_VER}+cu118.html"
else
  echo "Unsupported --cuda value: $CUDA_VARIANT"; exit 1
fi

pip install --index-url "$TORCH_INDEX_URL" \
  torch=="$TORCH_VER" torchvision=="$TV_VER" torchaudio=="$TA_VER"

pip install torch-geometric pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv \
  -f "$PYG_WHL_URL"

echo "[4/6] Installing FastAPI, Uvicorn and runtime deps"
# Pin numpy <2 to avoid compiled-extension ABI issues that often occur on servers
pip install "numpy<2" "fastapi>=0.110" "uvicorn[standard]>=0.30" "einops>=0.7" "requests>=2.32"

echo "[5/6] Verifying required assets"
ERR=0
if [[ ! -f "model_epoch_394_best_acc_0.9887.pth" ]]; then
  echo "Missing: model_epoch_394_best_acc_0.9887.pth"; ERR=1
fi
if [[ ! -f "dataset/train_dataset.pt" ]]; then
  echo "Missing: dataset/train_dataset.pt"; ERR=1
fi
if [[ ! -f "dataset/test_dataset.pt" ]]; then
  echo "Missing: dataset/test_dataset.pt"; ERR=1
fi
if [[ "$ERR" -ne 0 ]]; then
  echo "Some required files are missing. Please place them under project root/dataset and rerun."; exit 1
fi

echo "[6/6] Starting FastAPI server on $HOST:$PORT"
exec uvicorn model_api_fastapi:app --host "$HOST" --port "$PORT" --reload


